import torch
import numpy as np
import matplotlib.pyplot as plt


def plot_long_ecg(lead_data, rpeaks=None, tpeaks=None, leads_to_plot=None, display=True, save_to=''):
    """
    Plots specified ECG leads with optional R-peaks and T-peaks highlighted.
    - lead_data: PyTorch tensor of shape (n_leads, lead_length) with ECG waveforms.
    - rpeaks, tpeaks: torch.Tensor or list containing indices of peaks, applicable to all leads.
    - leads_to_plot: List of lead indices to plot. If None, plots the first lead only.
    - display: If True, displays the plot.
    - save_to: If provided, saves the plot to the given filepath.
    """

    if leads_to_plot is None:
        leads_to_plot = [0]  # Default to the first lead if not specified

    if lead_data.ndim == 1:
        lead_data = np.expand_dims(lead_data, axis=0)

    if not all(isinstance(lead, int) and 0 <= lead < lead_data.shape[0] for lead in leads_to_plot):
        raise ValueError("All lead indices must be valid integers within the range of lead_data")

    num_plots = len(leads_to_plot)
    fig, axes = plt.subplots(num_plots, 1, figsize=(22, 4 * num_plots))  # Adjusted for multiple plots

    if num_plots == 1:
        axes = [axes]  # Ensure axes is always a list for consistency

    for i, lead_idx in enumerate(leads_to_plot):
        lead = lead_data[lead_idx]

        # Plot the ECG lead
        axes[i].plot(lead, 'k-', linewidth=0.7, label='ECG')

        # Plot R-peaks if provided
        if rpeaks is not None:
            if isinstance(rpeaks, torch.Tensor):
                rpeaks = rpeaks.int().numpy()  # Convert tensor to numpy array if it's not already
            axes[i].plot(rpeaks, lead[rpeaks], 'rx', markersize=7, label='R-peaks')

        # Plot T-peaks if provided
        if tpeaks is not None:
            if isinstance(tpeaks, torch.Tensor):
                tpeaks = tpeaks.int().numpy()  # Convert tensor to numpy array if it's not already
            axes[i].plot(tpeaks, lead[tpeaks], 'go', markersize=7, label='T-peaks')

        axes[i].set_title(f'Lead {lead_idx + 1}')
        axes[i].set_ylabel('Amplitude')
        axes[i].set_xticks(np.arange(0, len(lead), 50))  # Adjust the range as per your data length
        axes[i].grid(True)
        axes[i].legend()

        # Rotate x-axis labels
        axes[i].tick_params(axis='x', rotation=45)

    fig.tight_layout()

    if save_to:
        plt.savefig(save_to)
    if display:
        plt.show()
    plt.close(fig)
